using Crux.Core.LangChain.Abstractions.Constants;
using Crux.Core.LangChain.Abstractions.Databases;
using Microsoft.Data.Sqlite;
using Newtonsoft.Json;

namespace Crux.Core.LangChain.Databases.Sqlite;
/// <summary>
/// 
/// </summary>
public class SqliteRagEngine : IDisposable, IVectorDatabase
{
    private readonly SqliteConnection _connection;
    private bool _disposed = false;
    /// <summary>
    /// 
    /// </summary>
    /// <param name="connectionString"></param>
    /// <param name="function"></param>
    /// <exception cref="ArgumentNullException"></exception>
    public SqliteRagEngine(
        string connectionString = ":memory:",
        string function = DistanceFunction.EuclideanDistance)
    {
        SQLitePCL.raw.SetProvider(new SQLitePCL.SQLite3Provider_e_sqlite3());

        if (string.IsNullOrEmpty(connectionString))
        {
            throw new ArgumentNullException(nameof(connectionString));
        }

        _connection = new SqliteConnection(connectionString);
        _connection.Open();

        Func<float[], float[], float> distanceFunction = function switch
        {
            DistanceFunction.EuclideanDistance => Utils.ComputeEuclideanDistance,
            DistanceFunction.ManhattanDistance => Utils.ComputeManhattanDistance,
            DistanceFunction.CosineSimilarity => Utils.CosineSimilarity,
            DistanceFunction.DotProductSimilarity => Utils.DotProductSimilarity,
            _ => Utils.ComputeEuclideanDistance,
        };

        _connection.CreateFunction(
                   "distance",
                   (string a, string b)
                       =>
                   {
                       var vecA = JsonConvert.DeserializeObject<float[]>(a, SqliteRagConstants.SerializerSettings);
                       var vecB = JsonConvert.DeserializeObject<float[]>(b, SqliteRagConstants.SerializerSettings);
                       if (vecA == null || vecB == null)
                           return 0f;

                       return distanceFunction(vecA, vecB);
                   });
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="connectionString"></param>
    /// <returns></returns>
    public static SqliteRagEngine Create(string connectionString = ":memory:")
        => new(connectionString);

    /// <inheritdoc />
    public async Task<IVectorCollection> GetCollectionAsync(string collectionName, int dimensions = 1536, CancellationToken cancellationToken = default)
    {
        if (_disposed)
        {
            throw new ObjectDisposedException(nameof(SqliteRagEngine));
        }

        if (!await IsTableExistsAsync(collectionName, cancellationToken).ConfigureAwait(false))
        {
            await CreateTableAsync(collectionName, cancellationToken).ConfigureAwait(false);
        }

        return new SqliteRagCollection(_connection, collectionName);
    }


    /// <inheritdoc />
    private async Task<bool> IsTableExistsAsync(string collectionName, CancellationToken cancellationToken = default)
    {
        if (_disposed)
        {
            throw new ObjectDisposedException(nameof(SqliteRagEngine));
        }

        using var command = _connection.CreateCommand();

        command.CommandText = $"SELECT name FROM sqlite_master WHERE type='table' AND name='{collectionName}'";
        var result = await command.ExecuteScalarAsync(cancellationToken).ConfigureAwait(false);

        return result != null;
    }

    /// <inheritdoc />
    private async Task CreateTableAsync(string collectionName, CancellationToken cancellationToken = default)
    {
        using var commandCreate = _connection.CreateCommand();
        commandCreate.CommandText = $"CREATE TABLE IF NOT EXISTS {collectionName} (id TEXT PRIMARY KEY, vector BLOB, document TEXT)";
        await commandCreate.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);

        using var commandIndex = _connection.CreateCommand();
        commandIndex.CommandText = $"CREATE INDEX idx_langchain_json_id ON {collectionName}(json_extract(document, '$.Metadata'))";
        await commandIndex.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc />
    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        _disposed = true;

        using var _ = _connection;

        GC.SuppressFinalize(this);
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="collectionName"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    public async Task DeleteCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
    {
        using var command = _connection.CreateCommand();
        command.CommandText = $"DROP TABLE IF EXISTS {collectionName}";

        await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
    }
}